import os

# from dotenv import load_dotenv



os.environ["OPENAI_API_KEY"] = "sk-ns79vLHscEaduFZA60C1A329895c4833Ab06B7D5B290A485"
os.environ["SERPAPI_API_KEY"] = "96c3d7030d5c5cc7ec30cd425e55b7d8e3fb5ffc8a44d9889c942710b2f3b1c6"
os.environ["OPENAI_API_BASE"]  ="https://aihubmix.com/v1"
os.environ["TAVILY_API_KEY"] = "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"

from llama_index.core import SimpleDirectoryReader , VectorStoreIndex
docs = SimpleDirectoryReader("docs").load_data()

index = VectorStoreIndex.from_documents(docs)

agent = index.as_query_engine()

resp1 = agent.query("工资应该是多少？")

resp2 = agent.query("9月份，刘磊的工资应该是多少？")


resp3 = agent.query("9月份，如果刘磊请了一天假，那么刘磊的工资应该是多少？")

print(resp1)
print(resp2)
print(resp3)

# 索引保存在本地
index.storage_context.persist()
